{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "from multiprocessing import cpu_count, shared_memory\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a large array of random values\n",
    "random = np.random.rand(400, 400, 1200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a shared memory array for the input and result arrays, copying the original input data\n",
    "shm_random = shared_memory.SharedMemory(create=True, size=random.data.nbytes)\n",
    "shared_random = np.ndarray(random.shape, dtype=random.dtype, buffer=shm_random.buf)\n",
    "shared_random[:, :, :] = random[:, :, :]\n",
    "shm_name_random = shm_random.name\n",
    "\n",
    "shm_result = shared_memory.SharedMemory(create=True, size=random.data.nbytes)\n",
    "shared_result = np.ndarray(random.shape, dtype=random.dtype, buffer=shm_result.buf)\n",
    "shm_name_result = shm_result.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_average(ary: np.ndarray) -> np.ndarray:\n",
    "    return ary + np.mean(ary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shm_add_average(\n",
    "    arguments: dict,\n",
    "):\n",
    "    existing_shm_input = shared_memory.SharedMemory(name=arguments[\"shm_name_input\"])\n",
    "    shm_ary_input = np.ndarray(\n",
    "        shape=arguments[\"shape\"],\n",
    "        dtype=arguments[\"dtype\"],\n",
    "        buffer=existing_shm_input.buf,\n",
    "    )\n",
    "    existing_shm_result = shared_memory.SharedMemory(name=arguments[\"shm_name_result\"])\n",
    "    shm_ary_result = np.ndarray(\n",
    "        shape=arguments[\"shape\"],\n",
    "        dtype=arguments[\"dtype\"],\n",
    "        buffer=existing_shm_result.buf,\n",
    "    )\n",
    "    shm_ary_result[arguments[\"lat_index\"], arguments[\"lon_index\"]] = add_average(\n",
    "        shm_ary_input[arguments[\"lat_index\"], arguments[\"lon_index\"]]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 23.9 s, sys: 2.69 s, total: 26.6 s\n",
      "Wall time: 1min 41s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "arguments_list = []\n",
    "for lat_index in range(shared_random.shape[0]):\n",
    "    for lon_index in range(shared_random.shape[1]):\n",
    "        arguments = {\n",
    "            \"lat_index\": lat_index,\n",
    "            \"lon_index\": lon_index,\n",
    "            \"shm_name_input\": shm_name_random,\n",
    "            \"shm_name_result\": shm_name_result,\n",
    "            \"dtype\": shared_random.dtype,\n",
    "            \"shape\": shared_random.shape,\n",
    "        }\n",
    "        arguments_list.append(arguments)\n",
    "\n",
    "# use a ProcessPoolExecutor to download the images in parallel\n",
    "with concurrent.futures.ProcessPoolExecutor(max_workers=cpu_count()) as executor:\n",
    "    # use the executor to map the download function to the iterable of arguments\n",
    "    executor.map(shm_add_average, arguments_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.61 s, sys: 69.9 ms, total: 1.68 s\n",
      "Wall time: 1.68 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "added_ary = random.copy()\n",
    "arguments_list = []\n",
    "for lat_index in range(random.shape[0]):\n",
    "    for lon_index in range(random.shape[1]):\n",
    "        added_ary[lat_index, lon_index] = add_average(random[lat_index, lon_index])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(shared_result, added_ary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(400, 400, 1200)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "shared_result.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(400, 400, 1200)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "added_ary.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
